package net.fruit.mq.client.impl.tracehook;

import java.util.ArrayList;

import org.apache.rocketmq.client.hook.SendMessageContext;
import org.apache.rocketmq.client.hook.SendMessageHook;
import org.apache.rocketmq.client.producer.SendStatus;
import org.apache.rocketmq.common.MixAll;

import net.fruit.mq.client.trace.common.MQTraceBean;
import net.fruit.mq.client.trace.common.MQTraceConstants;
import net.fruit.mq.client.trace.common.MQTraceContext;
import net.fruit.mq.client.trace.common.MQTraceType;
import net.fruit.mq.client.trace.dispatch.AsyncDispatcher;

public class MQClientSendMessageHookImpl implements SendMessageHook {
    /**
     * 该Hook该由哪个dispatcher发送轨迹数据
     */
    private AsyncDispatcher localDispatcher;


    public MQClientSendMessageHookImpl(AsyncDispatcher localDispatcher) {
        this.localDispatcher = localDispatcher;
    }


    @Override
    public String hookName() {
        return "MQClientSendMessageHook";
    }


    @Override
    public void sendMessageBefore(SendMessageContext context) {
        // 如果是消息轨迹本身的发送链路，则不需要再记录
        if (context == null || context.getMessage().getTopic().startsWith(MixAll.SYSTEM_TOPIC_PREFIX)) {
            return;
        }
        MQTraceContext mqContext = new MQTraceContext();
        mqContext.setTraceBeans(new ArrayList<MQTraceBean>(1));
        context.setMqTraceContext(mqContext);
        mqContext.setTraceType(MQTraceType.Pub);
        mqContext.setGroupName(context.getProducerGroup());
        MQTraceBean traceBean = new MQTraceBean();
        traceBean.setTopic(context.getMessage().getTopic());
        traceBean.setTags(context.getMessage().getTags());
        traceBean.setKeys(context.getMessage().getKeys());
        traceBean.setStoreHost(context.getBrokerAddr());
        traceBean.setBodyLength(context.getMessage().getBody().length);
        traceBean.setMsgType(context.getMsgType());
        mqContext.getTraceBeans().add(traceBean);
    }


    @Override
    public void sendMessageAfter(SendMessageContext context) {
        // 如果是消息轨迹本身的发送链路，则不需要再记录
        if (context == null || context.getMessage().getTopic().startsWith(MQTraceConstants.traceTopic) || context.getMqTraceContext() == null) {
            return;
        }
        if (context.getSendResult().getRegionId() == null
                || context.getSendResult().getRegionId().equals(MQTraceConstants.default_region)
                || context.getSendResult().isTraceOn()==false) {
            // if regionId is default or switch is false,skip it
            return;
        }
        MQTraceContext mqContext = (MQTraceContext) context.getMqTraceContext();
        MQTraceBean traceBean = mqContext.getTraceBeans().get(0);
        int costTime = (int) ((System.currentTimeMillis() - mqContext.getTimeStamp()) / mqContext.getTraceBeans().size());
        mqContext.setCostTime(costTime);
        if (context.getSendResult().getSendStatus().equals(SendStatus.SEND_OK)) {
            mqContext.setSuccess(true);
        } else {
            mqContext.setSuccess(false);
        }
        mqContext.setRegionId(context.getSendResult().getRegionId());
        traceBean.setMsgId(context.getSendResult().getMsgId());
        traceBean.setOffsetMsgId(context.getSendResult().getOffsetMsgId());
        traceBean.setStoreTime(mqContext.getTimeStamp() + costTime / 2);
        localDispatcher.append(mqContext);
    }
}
